# message_history/async_pg_history.py
import json
from typing import List, Optional
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import BaseMessage, messages_from_dict, messages_to_dict
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy import text
import config

# 创建全局异步引擎（推荐单例）
engine = create_async_engine(
    config.DATABASE_URL,
    echo=False,          # 生产环境设为 False
    pool_size=10,
    max_overflow=20,
    pool_pre_ping=True,
)

# 初始化表结构（可单独运行一次）
async def init_chat_history_table(table_name: str = "langchain_chat_messages"):
    async with engine.begin() as conn:
        await conn.execute(text(f"""
            CREATE TABLE IF NOT EXISTS {table_name} (
                id SERIAL PRIMARY KEY,
                session_id TEXT NOT NULL,
                message JSONB NOT NULL,
                created_at TIMESTAMPTZ DEFAULT NOW()
            );
        """))
        await conn.execute(text(f"""
            CREATE INDEX IF NOT EXISTS idx_session_id_{table_name} 
            ON {table_name} (session_id);
        """))


class AsyncPostgresChatMessageHistory(BaseChatMessageHistory):
    def __init__(self, session_id: str, table_name: str = "langchain_chat_messages"):
        if not session_id:
            raise ValueError("session_id must be provided")
        self.session_id = session_id
        self.table_name = table_name

    async def aget_messages(self) -> List[BaseMessage]:
        async with AsyncSession(engine) as session:
            result = await session.execute(
                text(f"""
                    SELECT message 
                    FROM {self.table_name} 
                    WHERE session_id = :session_id 
                    ORDER BY id ASC
                """),
                {"session_id": self.session_id}
            )
            rows = result.fetchall()
            # SQLAlchemy + asyncpg 自动将 JSONB 反序列化为 dict
            message_dicts = [row[0] for row in rows]  # row[0] is already a dict
            return messages_from_dict(message_dicts)

    async def aadd_messages(self, messages: List[BaseMessage]) -> None:
        if not messages:
            return
        async with AsyncSession(engine) as session:
            message_dicts = messages_to_dict(messages)
            for msg in message_dicts:
                # 🔧 关键修复：将 dict 序列化为 JSON 字符串
                await session.execute(
                    text(f"""
                        INSERT INTO {self.table_name} (session_id, message)
                        VALUES (:session_id, :message)
                    """),
                    {
                        "session_id": self.session_id,
                        "message": json.dumps(msg)  # ✅ 传 str 而非 dict
                    }
                )
            await session.commit()

    async def aclear(self) -> None:
        """异步清空当前 session 的所有消息"""
        async with AsyncSession(engine) as session:
            await session.execute(
                text(f"DELETE FROM {self.table_name} WHERE session_id = :session_id"),
                {"session_id": self.session_id}
            )
            await session.commit()

    # 同步方法：提示用户使用异步版本
    def get_messages(self) -> List[BaseMessage]:
        raise NotImplementedError("Use 'aget_messages()' in async context.")

    def add_messages(self, messages: List[BaseMessage]) -> None:
        raise NotImplementedError("Use 'aadd_messages()' in async context.")

    def clear(self) -> None:
        raise NotImplementedError("Use 'aclear()' in async context.")